'''
    project_name: YOLOV3代码复现
    file_neme: YOLOV3.py
    data: 2021年11月1日下午23:20:33
    author: 郑宇坤
'''
##导入所需包
import cv2
import numpy as np
# python绘图matplotlib
import matplotlib.pyplot as plt


# 可视化方法
def show(img):
    img_RGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    plt.imshow(img_RGB)
    plt.show()


# 导入YOLOv3模型
'''
    yolov3.weights YOLOV3-模型权重文件
    yolov3.cfg 定义了模型的结构与超参数
    coco.names 目标检测数据集的80个类别
'''
net = cv2.dnn.readNet('yolov3.weights', 'yolov3.cfg')

# 导入COCO数据集的类别
with open('coco.names', 'r') as f:
    classes = f.read().splitlines()
print(classes)

# 导入图象
img = cv2.imread('test_1.jpg')
show(img)
##print(img.shape)

# 获取图象的宽高
height, width, _ = img.shape
##print(height, width)

# 对图像进行预处理
'''
    1/255 对图像进行归一化
    (416, 416) 将图片缩放(输入YOLOV3的尺度必须是32的倍数)
    (0, 0, 0) 对每个像素进行减去像素 不操作
    swapRB = True 置换绿色通道与蓝色通道
    crop = False 不对图像进行裁剪
'''
blob = cv2.dnn.blobFromImage(img, 1 / 255, (416, 416), (0, 0, 0),
                             swapRB=True, crop=False)
# (1, 3, 416, 416) 一张图片 三通道RGB 大小416,416
# print(blob.shape)

# 将blob作为网络输入
net.setInput(blob)

# 获取三个尺度输出层的名称
layersNames = net.getLayerNames()
output_layers_names = [layersNames[i[0] - 1] for i in net.getUnconnectedOutLayers()]
# print(output_layers_names)

# 输入YOLOV3神经网络，前向推断预测
prediction = net.forward(output_layers_names)

'''
    (507, 85)
    13 * 13的尺度的 507个框(13 * 13 * 3)
    每一个框的85值 对应了x, y, w, h, c + COCO80个条件概率
'''
prediction[0].shape

# 26*26 尺度 (2028, 85)
prediction[1].shape

# 52*52 尺度 (8112, 85)
prediction[2].shape

# 查看第2个尺度 索引为99的 85维向量
# print(prediction[1][99])


# 构建空列表 将prediction输入结果解析到列表里

# 存放预测框坐标
boxes = []

# 存放置信度
objectness = []

# 存放类别概率
class_probs = []

# 存放预测框类别索引号
class_ids = []

# 存放预测框类别名称
class_names = []

for scale in prediction:  # 遍历三种尺度
    for bbox in scale:  # 遍历每个预测框
        obj = bbox[4]  # 获取该预测框的置信度confidence（objectness）
        class_scores = bbox[5:]  # 获取该预测框在COCO数据集80个类别的概率
        class_id = np.argmax(class_scores)  # 获取概率最高类别的索引号
        class_name = classes[class_id]  # 获取概率最高类别的名称
        class_prob = class_scores[class_id]  # 获取概率最高类别的概率

        # 获取预测框中心点坐标、预测框宽高
        center_x = int(bbox[0] * width)
        center_y = int(bbox[1] * height)
        w = int(bbox[2] * width)
        h = int(bbox[3] * height)
        # 计算预测框左上角坐标
        x = int(center_x - w / 2)
        y = int(center_y - h / 2)

        # 将每个预测框的结果存放至上面的列表中
        boxes.append([x, y, w, h])
        objectness.append(float(obj))
        class_ids.append(class_id)
        class_names.append(class_name)
        class_probs.append(class_prob)

# 结果均为10647 = (13 * 13 + 26 * 26 + 52 * 52) * 3
print(len(boxes))
print(len(objectness))

# 置信度预处理
# 将预测框置信度与各个类别的条件类别概率分别相乘,得到最终预测框中的置信度
# 结果仍然为10247
confidences = np.array(class_probs) * np.array(objectness)
# print(len(confidences))

# 后处理 置信度过滤,非极大值抑制NMS 剔除重复预测的框
# 指定置信度阈值,阈值越大,置信度过滤越强,小于0.1剔除
CONF_THRES = 0.4
# 指定NMS阈值,阈值越小,NMS越强
NMS_THRES = 0.4

# 传入数据各个框的坐标 置信度 两个阈值 进行置信度过滤和非极大值抑制
indexes = cv2.dnn.NMSBoxes(boxes, confidences, CONF_THRES, NMS_THRES)

# 随机给每个框一个颜色 进行可视化
colors = np.random.uniform(0, 255, size=(len(boxes), 3))
# colors = [[255,0,255],[0,0,255],[0,255,255],[0,255,0],[255,255,0],[255,0,0],[180,187,28],[223,155,6],
#          [94,218,121],[139,0,0],[77,169,10],[29,123,243],[66,77,229],[1,240,255],[140,47,240],[31,41,81],[29,123,243],[16,144,247],[151,57,224]]


# 遍历留下的每一个预测框，可视化
if len(indexes) == 0:
    print("没有检测到物体")

# 遍历留下的每一个预测框，可视化
if len(indexes) > 0:
    for i in indexes.flatten():
        # 获取坐标
        x, y, w, h = boxes[i]
        # 获取置信度
        confidence = str(round(confidences[i], 2))
        # 获取颜色，画框
        color = colors[i % len(colors)]
        cv2.rectangle(img, (x, y), (x + w, y + h), color, 8)

        # 写类别名称和置信度
        # 图片,添加的文字,左上角坐标,字体,字体大小,颜色,字体粗细
        string = '{} {}'.format(class_names[i], confidence)
        cv2.putText(img, string, (x, y + 20), cv2.FONT_HERSHEY_PLAIN, 3, (255, 255, 255), 5)

show(img)
cv2.imwrite('tested_1.jpg', img)